import torch
import torch.nn as nn
import numpy as np

# GPU or CPU
device = torch.device(
    "cuda") if torch.cuda.is_available() else torch.device("cpu")


class FCNN(nn.Module):
    def __init__(self, hidden_sizes, x_size, time_embed_size, condition_size, condition_embed_size):
        super(FCNN, self).__init__()
        input_size = x_size + time_embed_size + condition_embed_size
        bypass_size = time_embed_size + condition_embed_size

        self.layers = nn.ModuleList()
        self.bn_layers = nn.ModuleList()

        # First layer
        self.layers.append(nn.Linear(input_size, hidden_sizes[0]))
        self.bn_layers.append(nn.BatchNorm1d(hidden_sizes[0]))

        # Hidden layers
        for i in range(1, len(hidden_sizes)):
            self.layers.append(
                nn.Linear(hidden_sizes[i - 1] + bypass_size, hidden_sizes[i]))
            self.bn_layers.append(nn.BatchNorm1d(hidden_sizes[i]))

        # Output layer
        self.layers.append(nn.Linear(hidden_sizes[-1] + bypass_size, x_size))

        self.leaky_relu = nn.LeakyReLU(0.01)
        self.dropout = nn.Dropout(0.05)

        self.time_embeding = EmbedFC(1, time_embed_size)
        self.condition_embeding = EmbedFC(condition_size, condition_embed_size)

    def forward(self, x, time, condition, context_mask):
        time_embed = self.time_embeding(time)
        condition_embed = self.condition_embeding(
            condition) * (1.0 - context_mask)
        for i, layer in enumerate(self.layers[:-1]):
            x = torch.cat((x, time_embed, condition_embed), 1)
            x = layer(x)
            x = self.leaky_relu(x)
            x = self.bn_layers[i](x)
            x = self.dropout(x)

        x = torch.cat((x, time_embed, condition_embed), 1)
        x = self.layers[-1](x)

        return x

# A fully connected neural network for embed
class EmbedFC(nn.Module):
    def __init__(self, input_dim, emb_dim):
        super(EmbedFC, self).__init__()
        self.input_dim = input_dim
        layers = [
            nn.Linear(input_dim, emb_dim)
        ]
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        x = x.view(-1, self.input_dim)
        return self.model(x)

def ddpm_schedules(beta1, beta2, T):
    """
    Returns pre-computed schedules for DDPM sampling, training process.
    """
    assert beta1 < beta2 < 1.0, "beta1 and beta2 must be in (0, 1)"

    beta_t = (beta2 - beta1) * torch.arange(0, T +
                                            1, dtype=torch.float32) / T + beta1
    sqrt_beta_t = torch.sqrt(beta_t)
    alpha_t = 1 - beta_t
    log_alpha_t = torch.log(alpha_t)
    alphabar_t = torch.cumsum(log_alpha_t, dim=0).exp()

    sqrtab = torch.sqrt(alphabar_t)
    oneover_sqrta = 1 / torch.sqrt(alpha_t)

    sqrtmab = torch.sqrt(1 - alphabar_t)
    mab_over_sqrtmab_inv = (1 - alpha_t) / sqrtmab

    return {
        "alpha_t": alpha_t,  # \alpha_t
        "oneover_sqrta": oneover_sqrta,  # 1/\sqrt{\alpha_t}
        "sqrt_beta_t": sqrt_beta_t,  # \sqrt{\beta_t}
        "alphabar_t": alphabar_t,  # \bar{\alpha_t}
        "sqrtab": sqrtab,  # \sqrt{\bar{\alpha_t}}
        "sqrtmab": sqrtmab,  # \sqrt{1-\bar{\alpha_t}}
        # (1-\alpha_t)/\sqrt{1-\bar{\alpha_t}}
        "mab_over_sqrtmab": mab_over_sqrtmab_inv,
    }


class DDPM(nn.Module):
    def __init__(self, nn_model, betas, n_T, device, drop_prob=0.1):
        super(DDPM, self).__init__()
        self.nn_model = nn_model.to(device)
        num_params = sum(p.numel() for p in nn_model.parameters())
        print(f"Parameter number: {num_params*1e-6}M")
        # register_buffer allows accessing dictionary produced by ddpm_schedules
        # e.g. can access self.sqrtab later
        for k, v in ddpm_schedules(betas[0], betas[1], n_T).items():
            self.register_buffer(k, v)

        self.n_T = n_T
        self.device = device
        self.drop_prob = drop_prob
        self.loss_mse = nn.MSELoss()

    def forward(self, x, condition):
        """
        this method is used in training, so samples t and noise randomly
        """
        _ts = torch.randint(
            1, self.n_T+1, (x.shape[0],)).to(self.device)  # t ~ Uniform(0, n_T)
        noise = torch.randn_like(x)  # eps ~ N(0, 1)

        x_t = (
            self.sqrtab[_ts, None] * x
            + self.sqrtmab[_ts, None] * noise
        )  # This is the x_t, which is sqrt(alphabar) x_0 + sqrt(1-alphabar) * eps
        # We should predict the "error term" from this x_t. Loss is what we return.

        # dropout context with some probability
        context_mask = torch.bernoulli(
            torch.zeros_like(condition[:, 0:1])+self.drop_prob).to(self.device)

        # return MSE between added noise, and our predicted noise
        return self.loss_mse(noise, self.nn_model(x_t, _ts / self.n_T, condition, context_mask))

    def sample(self, n_sample, x_size, device, guide_w=0.0, condition=None):
        # we follow the guidance sampling scheme described in 'Classifier-Free Diffusion Guidance'
        # to make the fwd passes efficient, we concat two versions of the dataset,
        # one with context_mask=0 and the other context_mask=1
        # we then mix the outputs with the guidance scale, w
        # where w>0 means more guidance

        # x_T ~ N(0, 1), sample initial noise
        x_i = torch.randn(n_sample, x_size).to(device)
        condition = condition.unsqueeze(0).repeat(n_sample, 1).to(device)

        # don't drop context at test time
        context_mask = torch.zeros_like(condition[:, 0:1]).to(device)

        # double the batch
        condition = condition.repeat(2, 1)
        context_mask = context_mask.repeat(2, 1)
        context_mask[n_sample:] = 1.  # makes second half of batch context free
        x_i_store = []  # keep track of generated steps in case want to plot something
        for i in range(self.n_T, 0, -1):
            print(f'sampling timestep {i}\n')
            t_is = torch.tensor([i / self.n_T]).to(device)
            t_is = t_is.repeat(n_sample, 1)

            # double batch
            x_i = x_i.repeat(2, 1)
            t_is = t_is.repeat(2, 1)

            z = torch.randn(n_sample, x_size).to(device) if i > 1 else 0

            # split predictions and compute weighting
            eps = self.nn_model(x_i, t_is[:, 0], condition, context_mask)
            eps1 = eps[:n_sample]
            eps2 = eps[n_sample:]
            eps = (1.0+guide_w)*eps1 - guide_w*eps2
            x_i = x_i[:n_sample]

            x_i = (
                self.oneover_sqrta[i] * (x_i - eps * self.mab_over_sqrtmab[i])
                + self.sqrt_beta_t[i] * z
            )
            x_i_store.append(x_i.detach().cpu().numpy())
        x_i_store.reverse()
        x_i_store = np.array(x_i_store)
        x_i_store = torch.Tensor(x_i_store)
        return x_i, x_i_store
